package com.ciccwm.auth.dubbo.impl;

import com.ciccwm.auth.dubbo.OAuth2DubboService;
import com.ciccwm.auth.model.OAuth2TokenResponse;
import lombok.RequiredArgsConstructor;
import org.apache.dubbo.config.annotation.DubboService;
import org.springframework.security.authentication.UsernamePasswordAuthenticationToken;
import org.springframework.security.core.Authentication;
import org.springframework.security.core.GrantedAuthority;
import org.springframework.security.oauth2.core.*;
import org.springframework.security.oauth2.jwt.Jwt;
import org.springframework.security.oauth2.jwt.JwtEncoder;
import org.springframework.security.oauth2.server.authorization.OAuth2Authorization;
import org.springframework.security.oauth2.server.authorization.OAuth2AuthorizationService;
import org.springframework.security.oauth2.server.authorization.OAuth2TokenType;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2AccessTokenAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.authentication.OAuth2ClientAuthenticationToken;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClient;
import org.springframework.security.oauth2.server.authorization.client.RegisteredClientRepository;
import org.springframework.security.oauth2.server.authorization.token.DefaultOAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenContext;
import org.springframework.security.oauth2.server.authorization.token.OAuth2TokenGenerator;

import java.time.Instant;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * OAuth2 Dubbo服务实现类.
 */
@DubboService
@RequiredArgsConstructor
public class OAuth2DubboServiceImpl implements OAuth2DubboService {

    private final RegisteredClientRepository registeredClientRepository;
    private final OAuth2AuthorizationService authorizationService;
    private final OAuth2TokenGenerator<Jwt> tokenGenerator;

    @Override
    public OAuth2TokenResponse passwordGrant(String username, String password, String clientId, String clientSecret) {
        RegisteredClient registeredClient = registeredClientRepository.findByClientId(clientId);
        if (registeredClient == null) {
            throw new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodes.INVALID_CLIENT));
        }

        Authentication clientPrincipal = new OAuth2ClientAuthenticationToken(registeredClient,
            ClientAuthenticationMethod.CLIENT_SECRET_BASIC, clientSecret);

        Set<String> authorizedScopes = registeredClient.getScopes();

        Authentication principal = new UsernamePasswordAuthenticationToken(username, password);

        OAuth2TokenContext tokenContext = DefaultOAuth2TokenContext.builder()
            .registeredClient(registeredClient)
            .principal(principal)
            .authorizationGrantType(new AuthorizationGrantType("password"))
            .authorizedScopes(authorizedScopes)
            .authorizationGrant(principal)
            .build();

        OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
        if (generatedAccessToken == null) {
            throw new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR));
        }

        OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
            ((Jwt) generatedAccessToken).getTokenValue(),
            ((Jwt) generatedAccessToken).getIssuedAt(),
            ((Jwt) generatedAccessToken).getExpiresAt(),
            authorizedScopes.stream().collect(Collectors.toSet()));

        OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
            .principalName(principal.getName())
            .authorizationGrantType(new AuthorizationGrantType("password"))
            .accessToken(accessToken)
            .attribute("scope", authorizedScopes);

        OAuth2Authorization authorization = authorizationBuilder.build();
        this.authorizationService.save(authorization);

        return new OAuth2TokenResponse(
            accessToken.getTokenValue(),
            null,
            accessToken.getExpiresAt().getEpochSecond() - Instant.now().getEpochSecond(),
            String.join(" ", authorizedScopes)
        );
    }

    @Override
    public OAuth2TokenResponse clientCredentialsGrant(String clientId, String clientSecret) {
        RegisteredClient registeredClient = registeredClientRepository.findByClientId(clientId);
        if (registeredClient == null) {
            throw new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodes.INVALID_CLIENT));
        }

        Authentication clientPrincipal = new OAuth2ClientAuthenticationToken(registeredClient,
            ClientAuthenticationMethod.CLIENT_SECRET_BASIC, clientSecret);

        Set<String> authorizedScopes = registeredClient.getScopes();

        OAuth2TokenContext tokenContext = DefaultOAuth2TokenContext.builder()
            .registeredClient(registeredClient)
            .principal(clientPrincipal)
            .authorizationGrantType(new AuthorizationGrantType("client_credentials"))
            .authorizedScopes(authorizedScopes)
            .authorizationGrant(clientPrincipal)
            .build();

        OAuth2Token generatedAccessToken = this.tokenGenerator.generate(tokenContext);
        if (generatedAccessToken == null) {
            throw new OAuth2AuthenticationException(new OAuth2Error(OAuth2ErrorCodes.SERVER_ERROR));
        }

        OAuth2AccessToken accessToken = new OAuth2AccessToken(OAuth2AccessToken.TokenType.BEARER,
            ((Jwt) generatedAccessToken).getTokenValue(),
            ((Jwt) generatedAccessToken).getIssuedAt(),
            ((Jwt) generatedAccessToken).getExpiresAt(),
            authorizedScopes.stream().collect(Collectors.toSet()));

        OAuth2Authorization.Builder authorizationBuilder = OAuth2Authorization.withRegisteredClient(registeredClient)
            .principalName(clientPrincipal.getName())
            .authorizationGrantType(new AuthorizationGrantType("client_credentials"))
            .accessToken(accessToken)
            .attribute("scope", authorizedScopes);

        OAuth2Authorization authorization = authorizationBuilder.build();
        this.authorizationService.save(authorization);

        return new OAuth2TokenResponse(
            accessToken.getTokenValue(),
            null,
            accessToken.getExpiresAt().getEpochSecond() - Instant.now().getEpochSecond(),
            String.join(" ", authorizedScopes)
        );
    }

    @Override
    public String createAuthorizationCode(String clientId, String redirectUri) {
        // TODO: Implement authorization code creation
        throw new UnsupportedOperationException("Authorization code grant type not implemented yet");
    }

    @Override
    public OAuth2TokenResponse authorizationCodeGrant(String code, String clientId, String clientSecret, String redirectUri) {
        // TODO: Implement authorization code grant
        throw new UnsupportedOperationException("Authorization code grant type not implemented yet");
    }

    @Override
    public OAuth2TokenResponse implicitGrant(String clientId) {
        // TODO: Implement implicit grant
        throw new UnsupportedOperationException("Implicit grant type not implemented yet");
    }

    @Override
    public OAuth2TokenResponse refreshToken(String refreshToken, String clientId, String clientSecret) {
        // TODO: Implement refresh token grant
        throw new UnsupportedOperationException("Refresh token grant type not implemented yet");
    }
} 
